import requests
import pandas as pd
import json
import os
from datetime import datetime
import pymysql
from pandas import ExcelWriter
import openpyxl
from openpyxl.styles import Font, Alignment, PatternFill, Border, Side
from spiltExcel import capture_excel_sheets_screenshots

def get_api_data(url="http://127.0.0.1:6666/api/public/stock_rank_ljqs_ths"):
    """
    从API获取股票量价齐升数据
    参数:
        url: API地址，默认为量价齐升数据接口
    返回:
        处理后的JSON数据和股票代码列表
    """
    try:
        print(f"正在从API获取量价齐升数据: {url}")
        response = requests.get(url)
        response.raise_for_status()  # 检查请求是否成功
        
        data = response.json()
        print(f"成功获取量价齐升数据，共 {len(data)} 条记录")
        
        # 提取股票代码
        stock_codes = [item['股票代码'] for item in data if '股票代码' in item]
        print(f"提取到 {len(stock_codes)} 个股票代码")
        
        return data, stock_codes
    except Exception as e:
        print(f"获取量价齐升API数据失败: {e}")
        return [], []

def get_stock_details(stock_codes, db_config=None):
    """
    从数据库查询股票详细信息
    参数:
        stock_codes: 股票代码列表
        db_config: 数据库连接配置，如果为None则使用默认配置
    返回:
        字典，键为股票代码，值为对应的详细信息
    """
    if not stock_codes:
        print("没有股票代码，跳过数据库查询")
        return {}

    # 数据库连接参数，与zjlx.py保持一致
    default_db_config = {
        'host': '175.178.5.14',
        'port': 33066,
        'user': 'root',
        'password': '07M@Y0sqL1',
        'database': 'stock',
        'charset': 'utf8mb4',
        'connect_timeout': 10
    }

    # 如果提供了自定义配置，则使用自定义配置
    if db_config:
        for key, value in db_config.items():
            if value:  # 只更新非空值
                default_db_config[key] = value

    db_config = default_db_config

    # 构建SQL查询中的IN子句
    codes_str = ', '.join([f"'{code}'" for code in stock_codes])

    # 构建完整的SQL查询
    sql = f"""
    select
        spld.SYMBOL as '股票代码',
        spld.NAME as '名称',
        ROUND(spld.TOTAL_MARKET_CAP / 100000000, 2) as '市值',
        ROUND(spld.PE_DYNAMIC,2) as '市盈率(动)',
        ROUND(spld.YTD_CHG,2) as '年度涨幅(%)'
    from
        stock_personal_latest_data spld
        left join stock_sector_personal ssp on spld.ID=ssp.PERSONAL_ID
        left join stock_sector_last_data s on ssp.SECTOR_ID = s.ID
    where
        spld.SYMBOL in ({codes_str})
        and s.SECTOR_TYPE = '2'
    """

    print(f"执行SQL查询: {sql}")

    # 连接数据库并执行查询
    stock_details = {}
    try:
        print("尝试连接数据库...")
        conn = pymysql.connect(**db_config)
        print("数据库连接成功")

        with conn.cursor(pymysql.cursors.DictCursor) as cursor:
            print("执行SQL查询...")
            cursor.execute(sql)
            results = cursor.fetchall()
            print(f"查询结果: {len(results)} 条记录")
            
            for result in results:
                stock_code = result['股票代码']
                stock_details[stock_code] = result
                print(f"股票 {stock_code} 的详细信息已获取")
    except pymysql.err.OperationalError as e:
        print(f"数据库连接失败: {e}")
        print("将继续处理数据，但不包含详细信息")
    except Exception as e:
        print(f"数据库查询失败: {e}")
        print("将继续处理数据，但不包含详细信息")
    finally:
        if 'conn' in locals() and conn:
            conn.close()
            print("数据库连接已关闭")

    return stock_details

def merge_data(api_data, stock_details):
    """合并API数据和数据库查询到的详细信息"""
    merged_data = []
    
    for item in api_data:
        stock_code = item.get('股票代码')
        if stock_code and stock_code in stock_details:
            # 合并数据，优先保留API中的数据，然后添加数据库中的详细信息
            merged_item = {**item, **stock_details[stock_code]}
            merged_data.append(merged_item)
            print(f"合并股票 {stock_code} 的数据成功")
        else:
            # 如果没有详细信息，只保留API数据
            merged_data.append(item)
            print(f"股票 {stock_code} 没有详细信息，仅保留API数据")
    
    return merged_data

def save_to_excel(data, output_file=None):
    """将合并后的数据保存到Excel文件"""
    if not data:
        print("没有数据可保存")
        return None

    dateStamp = datetime.now().strftime("%Y%m%d")
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    output_dir = f"../spiltExcel/input/{dateStamp}"
    os.makedirs(output_dir, exist_ok=True)
    # 设置默认输出文件名
    if not output_file:
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        output_file = f"{output_dir}/ljqs_{timestamp}.xlsx"


    try:
        # 创建DataFrame并筛选需要的列
        df = pd.DataFrame(data)
        
        # 定义需要保留的列（按需求顺序排列）
        required_columns = [
            '序号', '名称', '市值', '市盈率(动)','最新价',
            '量价齐升天数', '累计换手率',
             '年度涨幅(%)', '所属行业'
        ]
        
        # 筛选并排序列
        available_columns = [col for col in required_columns if col in df.columns]
        df = df[available_columns]

        # 写入Excel文件并设置格式
        with ExcelWriter(output_file, engine='openpyxl') as writer:
            df.to_excel(writer, index=False, sheet_name='量价齐升股票')

            # 获取工作表
            worksheet = writer.sheets['量价齐升股票']

            # 设置列宽和行高
            for column in worksheet.columns:
                column_letter = openpyxl.utils.get_column_letter(column[0].column)
                worksheet.column_dimensions[column_letter].width = 18

            for i in range(1, worksheet.max_row + 1):
                worksheet.row_dimensions[i].height = 25

            # 设置字体和对齐方式
            font = Font(name='微软雅黑')
            alignment = Alignment(horizontal='center', vertical='center')

            # 创建黄色填充样式（用于表头）
            yellow_fill = PatternFill(start_color='FFFF00', end_color='FFFF00', fill_type='solid')

            # 创建边框样式
            thin_border = Border(
                left=Side(style='thin'),
                right=Side(style='thin'),
                top=Side(style='thin'),
                bottom=Side(style='thin')
            )

            # 应用格式到所有单元格
            for row_idx, row in enumerate(worksheet.iter_rows()):
                for cell in row:
                    cell.alignment = alignment
                    cell.border = thin_border
                    cell.font = font

                    # 为表头设置黄色背景
                    if row_idx == 0:
                        cell.fill = yellow_fill

        print(f"Excel文件已生成: {output_file}")
        return output_file
    except Exception as e:
        print(f"保存Excel文件失败: {e}")
        return None

def main():
    """主函数：获取API数据，查询数据库，合并数据并保存到Excel"""
    try:
        # 从API获取数据
        api_data, stock_codes = get_api_data()
        
        if not api_data:
            print("未获取到API数据，程序终止")
            return

        # 从数据库查询股票详细信息
        stock_details = get_stock_details(stock_codes)

        # 合并数据
        merged_data = merge_data(api_data, stock_details)

        # 保存到Excel
        excel_file = save_to_excel(merged_data)

        if excel_file:
            print(f"程序执行成功，Excel文件已保存至: {excel_file}")
        else:
            print("程序执行成功，但保存Excel文件失败")

    except Exception as e:
        print(f"程序执行出错: {e}")


def process():
    main()
if __name__ == "__main__":
    main()
